#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据库更新脚本 - 添加文本作业支持
"""

import os
import sys
from sqlalchemy import text

# 添加项目根目录到Python路径
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from exts import db
from app import app

def update_database():
    """更新数据库结构，添加文本作业支持"""
    with app.app_context():
        try:
            # 检查并添加homework表的homework_type字段
            print("正在检查homework表结构...")
            
            # 检查homework_type字段是否存在
            check_homework_type_sql = """
            SELECT COUNT(*) as count 
            FROM information_schema.COLUMNS 
            WHERE TABLE_SCHEMA = DATABASE() 
            AND TABLE_NAME = 'homework' 
            AND COLUMN_NAME = 'homework_type'
            """
            
            result = db.session.execute(text(check_homework_type_sql)).fetchone()
            
            if result[0] == 0:
                print("添加homework_type字段到homework表...")
                alter_homework_sql = """
                ALTER TABLE homework 
                ADD COLUMN homework_type VARCHAR(50) NOT NULL DEFAULT 'code' 
                COMMENT '作业类型：code或text'
                """
                db.session.execute(text(alter_homework_sql))
                print("✓ homework_type字段添加成功")
            else:
                print("✓ homework_type字段已存在")
                
            # 检查并添加uploadhomework表的新字段
            print("正在检查uploadhomework表结构...")
            
            # 检查text_content字段是否存在
            check_text_content_sql = """
            SELECT COUNT(*) as count 
            FROM information_schema.COLUMNS 
            WHERE TABLE_SCHEMA = DATABASE() 
            AND TABLE_NAME = 'uploadhomework' 
            AND COLUMN_NAME = 'text_content'
            """
            
            result = db.session.execute(text(check_text_content_sql)).fetchone()
            
            if result[0] == 0:
                print("添加text_content字段到uploadhomework表...")
                alter_uploadhomework_sql1 = """
                ALTER TABLE uploadhomework 
                ADD COLUMN text_content TEXT 
                COMMENT '文本作业内容'
                """
                db.session.execute(text(alter_uploadhomework_sql1))
                print("✓ text_content字段添加成功")
            else:
                print("✓ text_content字段已存在")
                
            # 检查submission_type字段是否存在
            check_submission_type_sql = """
            SELECT COUNT(*) as count 
            FROM information_schema.COLUMNS 
            WHERE TABLE_SCHEMA = DATABASE() 
            AND TABLE_NAME = 'uploadhomework' 
            AND COLUMN_NAME = 'submission_type'
            """
            
            result = db.session.execute(text(check_submission_type_sql)).fetchone()
            
            if result[0] == 0:
                print("添加submission_type字段到uploadhomework表...")
                alter_uploadhomework_sql2 = """
                ALTER TABLE uploadhomework 
                ADD COLUMN submission_type VARCHAR(50) NOT NULL DEFAULT 'code' 
                COMMENT '提交类型：code或text'
                """
                db.session.execute(text(alter_uploadhomework_sql2))
                print("✓ submission_type字段添加成功")
            else:
                print("✓ submission_type字段已存在")
            
            # 提交事务
            db.session.commit()
            print("\n数据库更新完成！")
            print("现在支持文本作业和代码作业两种类型了。")
            
        except Exception as e:
            db.session.rollback()
            print(f"数据库更新失败: {str(e)}")
            return False
            
    return True

if __name__ == "__main__":
    print("开始更新数据库...")
    if update_database():
        print("数据库更新成功！")
    else:
        print("数据库更新失败！")


